from ctypes import sizeof
from sys import flags
from scapy.all import *
from scapy.utils import PcapReader
import numpy as np
import os 
from os import path 
from txtop import TxtOperate


#遍历文件夹中所有文件
def scaner_file(url):
    file  = os.listdir(url)
    return file

#正则匹配文件名是否匹配
def choosefile(filename):
    if re.match(r'ips\w.+',filename) != None:
        return True
    else:
        return False

# 找到客户端IP，将第一个TCP包的源IP作为上行数据流的发起者（客户端），因为TCP流总是客户端先发起的
def find_port(packets):
    j = 0
    while(True):
        if packets[j].haslayer("TCP"):
            upsrc = packets[j].payload.src
            downsrc = packets[j].payload.dst
            print(upsrc,'   ',downsrc)
            break
        j = j+1

#基于流计算数据包长度熵
def calc_ent(x):
    """
        calculate shanno ent of x
    """
    x_value_list = set(x)
    ent = 0.0
    for x_value in x_value_list:
        p = float(x.count(x_value)) / len(x)
        logp = np.log2(p)
        ent -= p * logp
    return ent


def ReadIp(filepath):
    packets = rdpcap(filepath)
    j = 0
    while(True):
        if packets[j].haslayer("TCP"):
            upsrc = packets[j].payload.src
            downsrc = packets[j].payload.dst
            # print(upsrc,'   ',downsrc)
            break
        j = j+1
    return upsrc,downsrc


def split_hostfile(file_name,ipset):  #file_name为大的混杂流量pcap的文件路径， 
    packets = rdpcap(file_name)
    ip_num = len(ipset)
    l = len(packets)
    data_list = [[] for _ in range(ip_num)]
    for i in range(ip_num):     #最重要得到ip_num个主机的pcap
        for j in range(l):  # 遍历所有数据包
            if packets[j].haslayer("TCP"):
                ip1 = packets[j].payload.src  # 找出包含在ipset中ip对的第一个包，之后从ipset
                ip2 = packets[j].payload.dst  # 中将该ip对去掉，退出循环（该ip对被记录）
                if (ip1, ip2) in ipset:
                    src, dst = ip1, ip2
                    ipset.remove((src, dst))
                    break
                elif (ip2, ip1) in ipset:
                    src, dst = ip2, ip1
                    ipset.remove((src, dst))
                    break
                else:
                    continue
        for j in range(l):
            if packets[j].haslayer("TCP"):
                ip1 = packets[j].payload.src  # 再次遍历，将所有该ip对或ip逆对的包全部提取加入列表中
                ip2 = packets[j].payload.dst
                if (ip1, ip2) == (src, dst) or (ip2,ip1) == (src,dst):
                    data_list[i].append(packets[j])
    return data_list    


def save_file(dir, data_list):
    l = len(data_list)
    print('开始分主机保存pcap,共有'+str(l)+'个pcap包')
    for i in range(0,l):
        name = dir + 'hostfile'+str(i)+'.pcap'
        wrpcap(name, data_list[i])
    print('保存完成')    


if __name__ =="__main__":
    # url = ''
    x = 0
    y = 0
    ipset = set()
    url = os.path.abspath(os.path.dirname(__file__))
    file = scaner_file(url)

    dir_name = os.path.dirname(__file__) + '/'
    filepath = dir_name+'/first_step_retained.txt'
    txtort = TxtOperate()
    file = txtort.Readtxt(filepath)  #经过第一重过滤，剩下的文件名列表

    for filename in file:
        len_lst = []
        if choosefile(filename):
            x += 1
            packets = rdpcap(url+'/'+filename)
            for data in packets:
                if data.haslayer("TCP"):
                    len_lst.append(len(data))
            if calc_ent(len_lst) < 3.1:
                y += 1
                ip1,ip2 = ReadIp(dir_name+filename)
                if (ip1,ip2) not in ipset and (ip2,ip1) not in ipset:
                    ipset.add((ip1,ip2))
            # print(calc_ent(len_lst))
    # print(y)
    # print(float(y)/x)
            # with open(url+"/"+"https_length_entropy_true.txt", "a") as file:
            #     file.write(str(calc_ent(len_lst)) + "\n")
    data_list = split_hostfile(url+'/'+input('待输入的总的流量文件名:'),ipset)
    save_file(url+'/',data_list)
    print('Done!')